package algorithm;

import dataset.Dataset;
import scicalc.Calculate;

import java.util.ArrayList;
import java.util.List;

public class LVQ extends AbstractClusterAlgorithm {

    private double eta = 0.1;
    private int maxIter = 500;
    private int[] t = new int[]{0, 1, 1, 0, 0};
    private double[][] p;
    private int q = t.length;

    @Override
    public long fit(Dataset X) {
        long start = System.currentTimeMillis();
        this.p = new double[this.q][X.shape[1]];
        List<Integer> indices = new ArrayList<>();
        for (int i : this.t) {
            List<Integer> list = Calculate.where(X.label, i);
            while (true) {
                int newIndex = list.get(Calculate.randomInt(0, list.size()));
                if (!indices.contains(newIndex)) {
                    indices.add(newIndex);
                    break;
                }
            }
        }
        for (int i = 0; i < indices.size(); i++)
            this.p[i] = X.dataset[indices.get(i)];
        for (int time = 0; time < this.maxIter; time++) {
            int j = Calculate.randomInt(0, X.shape[0]);
            double[] distance = Calculate.sqrt(Calculate.sum(Calculate.pow(Calculate.minus(X.dataset[j], this.p), 2), Calculate.AXIS_1));
            int i = Calculate.argmin(distance);
            this.p[i] = this.t[i] == X.label[j] ?
                    Calculate.minus(this.p[i], Calculate.multi(Calculate.minus(this.p[i], X.dataset[j]), this.eta)) :
                    Calculate.minus(this.p[i], Calculate.multi(Calculate.minus(X.dataset[j], this.p[i]), this.eta));
        }
        this.labels = new int[X.shape[0]];
        for (int i = 0; i < X.shape[0]; i++) {
            double[] distance = Calculate.sqrt(Calculate.sum(Calculate.pow(Calculate.minus(X.dataset[i], this.p), 2), Calculate.AXIS_1));
            int argmin = Calculate.argmin(distance);
            this.labels[i] = argmin;
        }
        this.result.shape[0] = X.shape[0];
        this.result.shape[1] = X.shape[1];
        this.result.dataset = X.dataset;
        this.result.label = labels;
        return System.currentTimeMillis() - start;
    }

    public void setT(int... t) {
        this.t = t;
        this.q = t.length;
    }

    public void setEta(double eta) {
        this.eta = eta;
    }

    public void setMaxIter(int maxIter) {
        this.maxIter = maxIter;
    }
}
